from sklearn.manifold import TSNE
import umap
import pandas as pd
import plotly.express as px
from nicegui import ui
import warnings

# 忽略特定警告
warnings.filterwarnings("ignore", category=UserWarning, module="llvmlite")

# 载入数据
df = pd.read_csv('cleaned_personality_data.csv')

# 将布尔值转换为数值
df['Stage_fear'] = df['Stage_fear'].astype(int)
df['Drained_after_socializing'] = df['Drained_after_socializing'].astype(int)

# 分离特征和标签
features = df.drop(columns=['Personality'])
labels = df['Personality']

# 定义颜色映射
color_map = {'Introvert': 'blue', 'Extrovert': 'orange'}

# 使用 t-SNE 进行降维
tsne = TSNE(n_components=3, random_state=42)
df_tsne = tsne.fit_transform(features)

# 将 t-SNE 结果添加到 DataFrame
df_tsne = pd.DataFrame(df_tsne, columns=['TSNE1', 'TSNE2', 'TSNE3'])
df_tsne['Personality'] = labels

# 使用 UMAP 进行降维
umap_reducer = umap.UMAP(n_components=3)
df_umap = umap_reducer.fit_transform(features)

# 将 UMAP 结果添加到 DataFrame
df_umap = pd.DataFrame(df_umap, columns=['UMAP1', 'UMAP2', 'UMAP3'])
df_umap['Personality'] = labels

# UI布局
with ui.row().classes('p-4'):
    # 左侧筛选面板
    with ui.column().classes('w-1/4'):
        ui.label('筛选条件').classes('text-xl font-bold mb-2')

        # 性格类型筛选
        personality_select = ui.select(
            label='性格类型',
            options=['All'] + df['Personality'].unique().tolist(),
            value='All'
        )

        # 降维方法选择
        method_select = ui.select(
            label='降维方法',
            options=['t-SNE', 'UMAP'],
            value='t-SNE'
        )

        # 应用按钮
        def update():
            filtered = df_tsne if method_select.value == 't-SNE' else df_umap
            filtered = filtered if personality_select.value == 'All' else filtered[filtered['Personality'] == personality_select.value]
            update_charts(filtered)

        ui.button('应用筛选', on_click=update).classes('mt-2')

    # 右侧图表区
    with ui.column().classes('w-3/4'):
        # 初始3D散点图
        scatter3d_fig = px.scatter_3d(
            df_tsne,
            x='TSNE1',
            y='TSNE2',
            z='TSNE3',
            color='Personality',
            color_discrete_map=color_map,
            title='性格分析3D散点图（t-SNE降维）'
        )
        scatter3d_chart = ui.plotly(scatter3d_fig).style('height: 600px')

# 更新图表函数
def update_charts(filtered_df):
    if method_select.value == 't-SNE':
        scatter3d_fig = px.scatter_3d(
            filtered_df,
            x='TSNE1',
            y='TSNE2',
            z='TSNE3',
            color='Personality',
            color_discrete_map=color_map,
            title='性格分析3D散点图（t-SNE降维）'
        )
    else:
        scatter3d_fig = px.scatter_3d(
            filtered_df,
            x='UMAP1',
            y='UMAP2',
            z='UMAP3',
            color='Personality',
            color_discrete_map=color_map,
            title='性格分析3D散点图（UMAP降维）'
        )
    scatter3d_chart.update_figure(scatter3d_fig)

try:
    # 主程序
    # if __name__ in {"__main__", "__mp_main__"}:
    #     ui.run()
except KeyboardInterrupt:
    print("程序被用户中断")
    # 这里可以添加一些清理代码
    if __name__ in {"__main__", "__mp_main__"}:
        import sys

        port = 8083
        if '--port' in sys.argv:
            port_index = sys.argv.index('--port')
            if port_index + 1 < len(sys.argv):
                port = int(sys.argv[port_index + 1])
        ui.run(title='高级降维分析', port=port, show=True)